﻿import System.Environment
import Control.Monad 

-------------------------------------------------------------------------------
-- Frames 
-------------------------------------------------------------------------------

-- complementary nucleotids
compl :: Char -> Char
compl 'A' = 'T'
compl 'T' = 'A'
compl 'C' = 'G'
compl 'G' = 'C'

-- returns specified frame of sequence
frame :: Char -> Int -> [Char] -> [Char]
frame '+' 1 xs = xs
frame '+' 2 xs = tail xs
frame '+' 3 xs = tail $ frame '+' 2 xs
frame '-' 1 xs = reverse $ map compl xs
frame '-' 2 xs = tail $ frame '-' 1 xs
frame '-' 3 xs = tail $ frame '-' 2 xs

-- returns all frames of sequence
frames :: [Char] -> [[Char]]
frames xs = [frame '+' 1 xs, frame '+' 2 xs, frame '+' 3 xs, 
             frame '-' 1 xs, frame '-' 2 xs, frame '-' 3 xs]

-------------------------------------------------------------------------------
-- Codons
-------------------------------------------------------------------------------
data Codon = AAA | AAC | AAG | AAT |
             ACA | ACC | ACG | ACT |
             AGA | AGC | AGG | AGT |
             ATA | ATC | ATG | ATT |
             CAA | CAC | CAG | CAT |
             CCA | CCC | CCG | CCT |
             CGA | CGC | CGG | CGT |
             CTA | CTC | CTG | CTT |
             GAA | GAC | GAG | GAT |
             GCA | GCC | GCG | GCT |
             GGA | GGC | GGG | GGT |
             GTA | GTC | GTG | GTT |
             TAA | TAC | TAG | TAT |
             TCA | TCC | TCG | TCT |
             TGA | TGC | TGG | TGT |
             TTA | TTC | TTG | TTT 
             deriving (Eq, Show, Read)
             
readCodons :: [Char] -> [Codon]
readCodons [] = []
readCodons (_:[]) = []
readCodons (_:_:[]) = []
readCodons xs = read (take 3 xs) : readCodons (drop 3 xs)

-------------------------------------------------------------------------------
-- Aminoacids
-------------------------------------------------------------------------------
data Amino = Aln | Arg | Asn | Asp | Cys | Gln | Glu | Gly | His | Ile 
           | Leu | Lys | Met | Phe | Pro | Ser | Thr | Trp | Tyr | Val 
           | STP 
           deriving (Eq, Show, Read)
           
amino3 :: Codon -> Amino
amino3 c | elem c [GCC, GCT, GCA, GCG]           = Aln
         | elem c [CGT, CGC, CGA, CGG, AGA, AGG] = Arg
         | elem c [AAT, AAC]                     = Asn
         | elem c [GAT, GAC]                     = Asp
         | elem c [TGT, TGC]                     = Cys
         | elem c [CAA, CAG]                     = Gln
         | elem c [GAA, GAG]                     = Glu
         | elem c [GGT, GGC, GGA, GGG]           = Gly
         | elem c [CAT, CAC]                     = His
         | elem c [ATT, ATC, ATA]                = Ile
         | elem c [TTA, TTG, CTT, CTC, CTA, CTG] = Leu
         | elem c [AAA, AAG]                     = Lys
         | elem c [ATG]                          = Met
         | elem c [TTT, TTC]                     = Phe
         | elem c [CCT, CCC, CCA, CCG]           = Pro
         | elem c [TCT, TCC, TCA, TCG, AGT, AGC] = Ser
         | elem c [ACT, ACC, ACA, ACG]           = Thr
         | elem c [TGG]                          = Trp
         | elem c [TAT, TAC]                     = Tyr
         | elem c [GTT, GTC, GTA, GTG]           = Val
         | elem c [TAG, TGA, TAA]                = STP

amino1 :: Codon -> Char
amino1 c   | amino3 c == Aln = 'A'
           | amino3 c == Arg = 'R'
           | amino3 c == Asn = 'N'
           | amino3 c == Asp = 'D'
           | amino3 c == Cys = 'C'
           | amino3 c == Gln = 'Q'
           | amino3 c == Glu = 'E'
           | amino3 c == Gly = 'G'
           | amino3 c == His = 'H'
           | amino3 c == Ile = 'I'
           | amino3 c == Leu = 'L'
           | amino3 c == Lys = 'K'
           | amino3 c == Met = 'M'
           | amino3 c == Phe = 'F'
           | amino3 c == Pro = 'P'
           | amino3 c == Ser = 'S'
           | amino3 c == Thr = 'T'
           | amino3 c == Trp = 'W'
           | amino3 c == Tyr = 'Y'
           | amino3 c == Val = 'V'
           | amino3 c == STP = '*'
           
-------------------------------------------------------------------------------
-- Translation into chain of aminoacids
-------------------------------------------------------------------------------
-- list of one letter aminoacids
translate1 :: [Char] -> [Char]
translate1 ns = map amino1 (readCodons ns)

-- list of three letter aminoacids
translate3 :: [Char] -> [Amino]
translate3 ns = map amino3 (readCodons ns)

-------------------------------------------------------------------------------
-- ORFs
-------------------------------------------------------------------------------
startCodons = [ATG]
stopCodons = [TGA, TAA, TAG]

-- drop codons until it is found one of the specified list
dropUntilCodons :: [Codon] -> [Codon] -> [Codon]
dropUntilCodons [] _ = []
dropUntilCodons _ [] = []
dropUntilCodons ys (x:xs) | elem x ys = x:xs
                          | otherwise = dropUntilCodons ys xs
                          
-- take codons until it is found one of the specified list
takeUntilCodons :: [Codon] -> [Codon] -> [Codon]
takeUntilCodons [] cs = cs
takeUntilCodons _ [] = []
takeUntilCodons ys (x:xs)  | notElem x ys = x : takeUntilCodons ys xs
                           | otherwise    = [x]
                           
-- returns a list whith the orfs found in the sequence of codons
findOrf :: [Codon] -> [[Codon]]
findOrf [] = []
findOrf xs | not $ any (`elem` startCodons) xs = []
           | not $ any (`elem` stopCodons) xs = []
           | otherwise = orf:findOrf (tail orf)
           where 
             orf = takeUntilCodons stopCodons (dropUntilCodons startCodons xs)

-------------------------------------------------------------------------------
-- FASTA format
-------------------------------------------------------------------------------
data Fasta = Fasta {fastaHeader   :: [Char], 
                    fastaSequence :: [Char]
                   } deriving (Eq, Read, Show)

-- readFasta :: [Char] -> [Fasta]
                   
f = ">alelo1\nAAA\nTTT\n>alelo2\nCCC\nGGG"

readFastas :: [Char] -> [Fasta]
readFastas xs = map readFasta $ blocks xs

readFasta :: [Char] -> Fasta
readFasta xs = Fasta h (concat $ lines $ (tail s))
                 where (h, s) = break (== '\n') xs

blocks :: [Char] -> [[Char]]
blocks [] = []
blocks xs | rest == []  = [xs]
          | block == [] = blocks $ tail rest
          | rest /= []  = block : blocks (tail rest)
          where (block, rest) = break (=='>') xs

-------------------------------------------------------------------------------
-- Main
-------------------------------------------------------------------------------
alelo = "TCAAGATGCAATTATGCCAATTCTCCACCACATGCCTATCCGATAAATAA"
alelo0 = "ATGGCAAGAATAGCTTGCGAAAACGCACTCGCAACGGCAAGCTAATTAG"
alelo6 = "ATGGCCAGAATAGCTTGCGAAAATGCACTCGCAACAGCAAGCTAATTAG"

main :: IO ()
main = do
  putStrLn "nhbiomod 0.2"
--  putStrLn "Escribe una cadena de ADN y pulsa [INTRO]:"
--  f <- return alelo
  fs <- getContents  
  mapM_ reportFasta (readFastas fs)
  
--  mapM_ reports (readFastas fs)

reportFasta (Fasta h s) = do
    putStrLn h
    putStrLn "-----------------------------------------------------------------"
    reports s

reports xs = mapM_ report $ zip ["+1", "+2", "+3", "-1", "-2", "-3"] (frames $ xs)

report (o, fn) = do      
      putStrLn (o ++ ":" ++ fn)
      putStrLn " - Codons: "
      putStrLn $ "   " ++ (show codons)
      putStrLn " - Translation:"
      putStr   "   "
      putStrLn $ translate1 fn
      putStr   "   "
      print $ translate3 fn
      putStrLn $ " - Orfs: " ++ (show $ length orfs)
      mapM (putStrLn . ("   " ++) . show) orfs
      putStrLn ""
      where 
        codons = readCodons fn
        orfs = findOrf $ codons
 
